package filter

import (
	"encoding/base64"
	"encoding/json"
	"fmt"
	"gin-fileserver/pkg/setting"
	"strings"

	log "github.com/sirupsen/logrus"

	"github.com/Baidu-AIP/golang-sdk/aip/censor"
)

type AIFilterResult struct {
	LogID      int64          `json:"log_id"`
	Conclusion string         `json:"conclusion"`
	Data       []AIFilterData `json:"data"`
}

type AIFilterData struct {
	Msg         string  `json:"msg"`
	Probability float64 `json:"probability"`
	Type        int64   `json:"type"`
	Words       *string `json:"words,omitempty"`
}

var client *censor.ContentCensorClient

const useAI = true

func Setup() {
	censor.NewClient(setting.Cfg.AppSetting.BaiduAIAppKey, setting.Cfg.AppSetting.BaiduAISecretKey)
	client = censor.NewClient(setting.Cfg.AppSetting.BaiduAIAppKey, setting.Cfg.AppSetting.BaiduAISecretKey)
}

func Validate(bytes []byte, gif bool) error {
	if !useAI {
		return nil
	}
	base64d := base64.StdEncoding.EncodeToString(bytes)
	g := 0
	if gif {
		g = 1
	}
	res := client.ImgCensor(base64d, map[string]interface{}{"imgType": g})
	var r AIFilterResult
	if err := json.Unmarshal([]byte(res), &r); err != nil {
		log.Errorln(err.Error())
		return nil
	}
	if len(r.Data) > 0 {
		var reasons []string
		for _, d := range r.Data {
			reasons = append(reasons, d.Msg)
		}
		return fmt.Errorf(strings.Join(reasons, "|"))
	}
	return nil
}
